#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <assert.h>
#include <mpi.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_sort_long.h>

#include "halofinder.h"

void getDomainOrigin_shiftpartshalos(float domainOrigin[3])
{
  domainOrigin[0] = (float) (haloFinderData.BoxLengthSnapshot/((float) (NumDomains[0]))*((float) (domainCoords[0])) - haloFinderData.DomainBuffSize);
  domainOrigin[1] = (float) (haloFinderData.BoxLengthSnapshot/((float) (NumDomains[1]))*((float) (domainCoords[1])) - haloFinderData.DomainBuffSize);
  domainOrigin[2] = (float) (haloFinderData.BoxLengthSnapshot/((float) (NumDomains[2]))*((float) (domainCoords[2])) - haloFinderData.DomainBuffSize);

#ifndef LIGHTCONE
  wrappos(&(domainOrigin[0]),haloFinderData.BoxLengthSnapshot);
  wrappos(&(domainOrigin[1]),haloFinderData.BoxLengthSnapshot);
  wrappos(&(domainOrigin[2]),haloFinderData.BoxLengthSnapshot);
#endif
}

void point_global2domain(float *px, float *py, float *pz, float domainOrigin[3])
{
  px[0] -= domainOrigin[0];
  py[0] -= domainOrigin[1];
  pz[0] -= domainOrigin[2];

#ifndef LIGHTCONE
  wrappos(px,haloFinderData.BoxLengthSnapshot);
  wrappos(py,haloFinderData.BoxLengthSnapshot);
  wrappos(pz,haloFinderData.BoxLengthSnapshot);
#endif
}

void point_domain2global(float *px, float *py, float *pz, float domainOrigin[3])
{
  px[0] += domainOrigin[0];
  py[0] += domainOrigin[1];
  pz[0] += domainOrigin[2];

#ifndef LIGHTCONE
  wrappos(px,haloFinderData.BoxLengthSnapshot);
  wrappos(py,haloFinderData.BoxLengthSnapshot);
  wrappos(pz,haloFinderData.BoxLengthSnapshot);
#endif
}

void shiftpartshalosglobal2domain(void)
{
  float domainOrigin[3];
  long i;
  
#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"shifting parts/halos global to domain\n");
#endif
  
  getDomainOrigin_shiftpartshalos(domainOrigin);
  
#ifdef DEBUG
#if DEBUG_LEVEL > 2
  fprintf(stderr,"%d: domainOrigin = %f|%f|%f\n",ThisTask,domainOrigin[0],domainOrigin[1],domainOrigin[2]);
#endif
#endif

#ifdef USEOPENMP
#pragma omp parallel for schedule(static) default(none) private(i) shared(NumParts,PartsPx,PartsPy,PartsPz,domainOrigin,ThisTask,stderr,haloFinderData)  
#endif
  for(i=0;i<NumParts;++i)
    {

#ifdef DEBUG
      if(DEBUG_LEVEL > 3)
	if(i%(NumParts/100) == 0)
	  fprintf(stderr,"%d: %ld of %ld (%0.2f percent)\n",ThisTask,i,NumParts,((double) i)/((double) NumParts)*100.0);
#endif
      
      point_global2domain(PartsPx+i,PartsPy+i,PartsPz+i,domainOrigin);
    }

#ifdef DEBUG
  if(ThisTask == 0 && DEBUG_LEVEL > 0)
    fprintf(stderr,"done shifting parts global to domain\n");
#endif  

#ifdef USEOPENMP
#pragma omp parallel for schedule(static) default(none) private(i) shared(NumHalos,Halos,domainOrigin,ThisTask,stderr,haloFinderData)  
#endif
  for(i=0;i<NumHalos;++i)
    {
      point_global2domain(Halos[i].pos+0,Halos[i].pos+1,Halos[i].pos+2,domainOrigin);
    }
  
#ifdef DEBUG
  if(ThisTask == 0 && DEBUG_LEVEL > 0)
    fprintf(stderr,"done shifting halos global to domain\n");
#endif
}

void shifthalosglobal2domain(void)
{
  float domainOrigin[3];
  long i;
  
#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"shifting halos global to domain\n");
#endif
  
  getDomainOrigin_shiftpartshalos(domainOrigin);
  
#ifdef DEBUG
#if DEBUG_LEVEL > 2
  fprintf(stderr,"%d: domainOrigin = %f|%f|%f\n",ThisTask,domainOrigin[0],domainOrigin[1],domainOrigin[2]);
#endif
#endif
  
#ifdef USEOPENMP
#pragma omp parallel for schedule(static) default(none) private(i) shared(NumHalos,Halos,domainOrigin,ThisTask,stderr,haloFinderData)  
#endif
  for(i=0;i<NumHalos;++i)
    {
      point_global2domain(Halos[i].pos+0,Halos[i].pos+1,Halos[i].pos+2,domainOrigin);
    }
  
#ifdef DEBUG
  if(ThisTask == 0 && DEBUG_LEVEL > 0)
    fprintf(stderr,"done shifting halos global to domain\n");
#endif
}

void shiftpartshalosdomain2global(void)
{
  float domainOrigin[3];
  long i;
  
#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"shifting parts/halos domain to global\n");
#endif

  getDomainOrigin_shiftpartshalos(domainOrigin);
  
#ifdef DEBUG
#if DEBUG_LEVEL > 2
  fprintf(stderr,"%d: domainOrigin = %f|%f|%f\n",ThisTask,domainOrigin[0],domainOrigin[1],domainOrigin[2]);
#endif
#endif
  
  #ifdef USEOPENMP
#pragma omp parallel for schedule(static) default(none) private(i) shared(NumParts,PartsPx,PartsPy,PartsPz,domainOrigin,ThisTask,stderr,haloFinderData)
#endif
  for(i=0;i<NumParts;++i)
    {

#ifdef DEBUG
      if(DEBUG_LEVEL > 3)
        if(i%(NumParts/100) == 0)
          fprintf(stderr,"%d: %ld of %ld (%0.2f percent)\n",ThisTask,i,NumParts,((double) i)/((double) NumParts)*100.0);
#endif

      point_domain2global(PartsPx+i,PartsPy+i,PartsPz+i,domainOrigin);
    }

#ifdef DEBUG
  if(ThisTask == 0 && DEBUG_LEVEL > 0)
    fprintf(stderr,"done shifting parts domain to global\n");
#endif

#ifdef USEOPENMP
#pragma omp parallel for schedule(static) default(none) private(i) shared(NumHalos,Halos,domainOrigin,ThisTask,stderr,haloFinderData)
#endif
  for(i=0;i<NumHalos;++i)
    {
      point_domain2global(Halos[i].pos+0,Halos[i].pos+1,Halos[i].pos+2,domainOrigin);
    }

#ifdef DEBUG
  if(ThisTask == 0 && DEBUG_LEVEL > 0)
    fprintf(stderr,"done shifting halos domain to global\n");
#endif
}

void shifthalosdomain2global(void)
{
  float domainOrigin[3];
  long i;
  
#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"shifting halos domain to global\n");
#endif

  getDomainOrigin_shiftpartshalos(domainOrigin);
  
#ifdef DEBUG
#if DEBUG_LEVEL > 2
  fprintf(stderr,"%d: domainOrigin = %f|%f|%f\n",ThisTask,domainOrigin[0],domainOrigin[1],domainOrigin[2]);
#endif
#endif

  #ifdef USEOPENMP
#pragma omp parallel for schedule(static) default(none) private(i) shared(NumHalos,Halos,domainOrigin,ThisTask,stderr,haloFinderData)
#endif
  for(i=0;i<NumHalos;++i)
    {
      point_domain2global(Halos[i].pos+0,Halos[i].pos+1,Halos[i].pos+2,domainOrigin);
    }

#ifdef DEBUG
  if(ThisTask == 0 && DEBUG_LEVEL > 0)
    fprintf(stderr,"done shifting halos domain to global\n");
#endif
}

void scatterpeakshalos(void)
{
  int recvcount,*displs,*sendcounts;
  long i,j,NumBuffHalos;
  Halo *tmpHalos,*buffHalos;
  size_t *sindex;
  int haloDestTask,coords[3];
  long *rank;
  float dx,dy,dz;
  Halo sourceHalo,saveHalo;
  long rankSource,rankSave,dest;
  
#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"doing scatter of halos\n");
#endif

  dx = (float) (haloFinderData.BoxLengthSnapshot/NumDomains[0]);
  dy = (float) (haloFinderData.BoxLengthSnapshot/NumDomains[1]);
  dz = (float) (haloFinderData.BoxLengthSnapshot/NumDomains[2]);
  
  if(ThisTask == 0)
    {
      //move Halos to buffHalos
      buffHalos = Halos;
      Halos = NULL;
      NumBuffHalos = NumHalos;
      NumHalos = 0;
      
      //get dest tasks
      rank = (long*)malloc(sizeof(long)*NumBuffHalos);
      assert(rank != NULL);
      for(i=0;i<NumBuffHalos;++i)
        {
          coords[0] = buffHalos[i].pos[0]/dx;
          coords[1] = buffHalos[i].pos[1]/dy;
          coords[2] = buffHalos[i].pos[2]/dz;
          for(j=0;j<3;++j)
            {
#ifndef LIGHTCONE
              while(coords[j] < 0)
                coords[j] += NumDomains[j];
              while(coords[j] >= NumDomains[j])
                coords[j] -= NumDomains[j];
#else
	      if(coords[j] < 0)
                coords[j] = 0;
              if(coords[j] >= NumDomains[j])
                coords[j] = NumDomains[j]-1;
#endif
            }
          
          MPI_Cart_rank(cartComm,coords,&haloDestTask);
	  rank[i] = (long) (haloDestTask);
        }
      
      //get scatter parms
      sendcounts = (int*)malloc(sizeof(int)*NTasks);
      assert(sendcounts != NULL);
      displs = (int*)malloc(sizeof(int)*NTasks);
      assert(displs != NULL);
      for(i=0;i<NTasks;++i)
        sendcounts[i] = 0;
      for(i=0;i<NumBuffHalos;++i)
        sendcounts[rank[i]] += 1;
      displs[0] = 0;
      for(i=1;i<NTasks;++i)
        displs[i] = displs[i-1]+sendcounts[i-1];
      
#ifdef DEBUG
      if(ThisTask == 0 && DEBUG_LEVEL > 1)
	{
	  for(i=0;i<NTasks;++i)
	    fprintf(stderr,"%ld: sendcounts,displs = %d|%d\n",i,sendcounts[i],displs[i]);
 	}
#endif
      
      //sort halos
      sindex = (size_t*)malloc(sizeof(size_t)*NumBuffHalos);
      assert(sindex != NULL);
      gsl_sort_long_index(sindex,rank,(size_t) 1,(size_t) NumBuffHalos);
      
      for(i=0;i<NumBuffHalos;++i)
	rank[sindex[i]] = i;
      free(sindex);
      
      for(i=0;i<NumBuffHalos;++i) /* reoder with an in-place algorithm - see Gadget-2 for details - destroys rank */
	{
	  if(i != rank[i])
	    {
	      sourceHalo = buffHalos[i];
	      rankSource = rank[i];
	      dest = rank[i];
                      
	      do
		{
		  saveHalo = buffHalos[dest];
		  rankSave = rank[dest];
                          
                  buffHalos[dest] = sourceHalo;
		  rank[dest] = rankSource;
                  
		  if(dest == i)
		    break;
		  
		  sourceHalo = saveHalo;
		  rankSource = rankSave;
                  
		  dest = rankSource;
		}
	      while(1);
	    }
	}
      
      free(rank);
    }
  
  //send each task the number of halos it will recv
  MPI_Scatter(sendcounts,1,MPI_INT,&recvcount,1,MPI_INT,0,MPI_COMM_WORLD);
  if(recvcount > 0)
    {
      tmpHalos = (Halo*)realloc(Halos,sizeof(Halo)*(NumHalos+recvcount));
      assert(tmpHalos != NULL);
      Halos = tmpHalos;
    }
  
  //send the Halos
  MPI_Scatterv(buffHalos,sendcounts,displs,MPI_Halo,
               Halos+NumHalos,recvcount,MPI_Halo,
               0,MPI_COMM_WORLD);
  NumHalos += recvcount;
  
#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"did scatter of halos\n");
#endif
  
  if(ThisTask == 0)
    {
      free(buffHalos);
      free(sendcounts);
      free(displs);
    }
}

void gatherpeakshalos(void)
{
  int *recvcounts,*displs,sendcount;
  long i,totRecv;
  Halo *buffHalos;
    
  //get recvcounts and displs
  if(ThisTask == 0)
    {
      recvcounts = (int*)malloc(sizeof(int)*NTasks);
      assert(recvcounts != NULL);
      displs = (int*)malloc(sizeof(int)*NTasks);
      assert(displs != NULL);
    }
  sendcount = NumHalos;
  MPI_Gather(&sendcount,1,MPI_INT,
             recvcounts,1,MPI_INT,0,
             MPI_COMM_WORLD);
  if(ThisTask == 0)
    {
      displs[0] = 0;
      totRecv = recvcounts[0];
      for(i=1;i<NTasks;++i)
        {
          displs[i] = recvcounts[i-1] + displs[i-1];
          totRecv += recvcounts[i];
        }
    }
  
#ifdef DEBUG
  if(ThisTask == 0 && DEBUG_LEVEL > 1)
    {
      fprintf(stderr,"totRecv = %ld\n",totRecv);
      for(i=0;i<NTasks;++i)
        fprintf(stderr,"%ld: recvcounts,displs = %d|%d\n",i,recvcounts[i],displs[i]);
    }
#endif
  
  //do gather of halos
  if(ThisTask == 0)
    {
      buffHalos = (Halo*)malloc(sizeof(Halo)*totRecv);
      assert(buffHalos != NULL);
    }
  MPI_Gatherv(Halos,sendcount,MPI_Halo,
              buffHalos,recvcounts,displs,MPI_Halo,
              0,MPI_COMM_WORLD);
  destroyhalos();
  if(ThisTask == 0)
    {
      NumHalos = totRecv;
      Halos = buffHalos;
    }

#ifdef DEBUG
  if(ThisTask == 0)
    fprintf(stderr,"did gather of halos\n");
#endif

  if(ThisTask == 0)
    {
      free(recvcounts);
      free(displs);
    }
}

void setupdomaindecomp(void)
{
  int reorderTasks,periods[3];
  double dx,dy,dz;
  
  //get # of dims in each direction from MPI_Dims_create
  NumDomains[0] = 0;
  NumDomains[1] = 0;
  NumDomains[2] = 0;
  MPI_Dims_create(NTasks,3,NumDomains);
  assert((NumDomains[0] >= 2 && NumDomains[1] >= 2 && NumDomains[2] >= 2)
	 || 
	 (NumDomains[0] == 1 && NumDomains[1] == 1 && NumDomains[2] == 1));
  
  //make cart topology using MPI
  reorderTasks = 0;
#ifdef LIGHTCONE
  periods[0] = 0;
  periods[1] = 0;
  periods[2] = 0;
#else
  periods[0] = 1;
  periods[1] = 1;
  periods[2] = 1;
#endif
  MPI_Cart_create(MPI_COMM_WORLD,3,NumDomains,periods,reorderTasks,&cartComm);
  
  //get coords
  MPI_Cart_coords(cartComm,ThisTask,3,domainCoords);
  
  //get domain center
  dx = haloFinderData.BoxLengthSnapshot/NumDomains[0];
  dy = haloFinderData.BoxLengthSnapshot/NumDomains[1];
  dz = haloFinderData.BoxLengthSnapshot/NumDomains[2];
  
#ifdef DEBUG
  if(DEBUG_LEVEL > 1)
    {
      fprintf(stderr,"%2d: NumDomains = %d|%d|%d, domainCoords = %d|%d|%d, sizes = %lf|%lf|%lf\n",ThisTask,
	      NumDomains[0],NumDomains[1],NumDomains[2],domainCoords[0],domainCoords[1],domainCoords[2],
	      dx,dy,dz);
    }
#endif
  
  /* Halo data typedef
     typedef struct {
     float pos[3];
     float vel[3]; 
     long index;
     float peakOverdens; 
     float radius;
     float vcirc;
     } Halo;
  */
  Halo tmpHalo;
  int i,j;
  MPI_Datatype type[8] = {MPI_LB,MPI_FLOAT,MPI_FLOAT,MPI_LONG,MPI_FLOAT,MPI_FLOAT,MPI_FLOAT,MPI_UB};
  int blocklen[8] = {1,3,3,1,1,1,1,1};
  MPI_Aint disp[8];

  //make an MPI type for Halo
  i = 0;
  MPI_Get_address(&(tmpHalo),disp+i);
  ++i;
  MPI_Get_address(&(tmpHalo.pos[0]),disp+i);
  ++i;
  MPI_Get_address(&(tmpHalo.vel[0]),disp+i);
  ++i;
  MPI_Get_address(&(tmpHalo.index),disp+i);
  ++i;
  MPI_Get_address(&(tmpHalo.mass),disp+i);
  ++i;
  MPI_Get_address(&(tmpHalo.radius),disp+i);
  ++i;
  MPI_Get_address(&(tmpHalo.vcirc),disp+i);
  ++i;
  disp[i] = disp[0] + sizeof(Halo);
  ++i;
  for(j=1;j<i;++j)
    disp[j] -= disp[0];
  disp[0] -= disp[0];
  MPI_Type_create_struct((int) i,blocklen,disp,type,&MPI_Halo);
  MPI_Type_commit(&MPI_Halo);
}

void destroydomaindecomp(void)
{
  MPI_Comm_free(&cartComm);
  MPI_Type_free(&MPI_Halo);
}
